import warnings

import torch
import torch.nn as nn

import torch.nn.functional as F
from mmcv.runner import BaseModule

INF = 1e8

import torch
from torch import nn


from mmdet.core import multi_apply

class instancenorm_head(BaseModule):
    def __init__(self,init_cfg=None):
        super(instancenorm_head, self).__init__(init_cfg)
    def forward(self, feats):
        # feats = list(feats)
        # feats1 = []
        # for feat in feats:
        #     h, w = feat.shape[1], feat.shape[2]
        #     mean_value = nn.functional.adaptive_avg_pool1d(feat, (1,1))
        #     mean_value = F.upsample(input=mean_value, size=(h, w), mode='bilinear')
        #     feat = feat - mean_value
        #     feats1.append(feat)
        feats = list(feats)
        h, w = feats[0].shape[2], feats[0].shape[3]
        mean_value = nn.functional.adaptive_avg_pool2d(feats[0], 1)
        mean_value = F.upsample(input=mean_value, size=(h, w), mode='bilinear')
        feats[0] = feats[0] - mean_value
        return feats[0]
        # return multi_apply(self.forward_single, feats)

    def forward_single(self, x):
        reid_feat = x
        return reid_feat